# coding=UTF-8
import random
print("start to build train data pair!")
from const_ import *
#with open("data",encoding="utf-8") as f:
with open("./data.txt") as f:
    lines = f.readlines()

relation_dict = {}

def if_e1_front_e2(q,a):
    entity2mask,e1_front_e2=find_entity1and2(q,a)
    return e1_front_e2

train_data = []

for i,line in enumerate(lines):
    if i%3==0:
        e1_front_e2 = if_e1_front_e2(extract_question(line),lines[i+1])
        one_pair = []
        one_pair.append(extract_id(lines[i+1]))
        one_pair.append(extract_question(line))
        one_pair.append(e1_front_e2)
        train_data.append(one_pair)


print("start to load word dict!")


word_dict = {}
with open("word_dict.txt","r") as f:
    lines2 = f.readlines()
    for line in lines2:
        temp = line.split(" ")
        word_dict[temp[0]] = int(temp[1].replace("\n",""))

for item in train_data:
        new_sentence = replace_entity_in_logic_form(item[1], lines[int(item[0]) * 3 - 2])
        new_sentence = replace_number_and_date_after_entity_removed(new_sentence)
        item[1] = new_sentence

# for item in train_data:
#     entity2mask,e1_front_e2 = find_entity1and2(item[1],lines[int(item[0])*3-2])
#     if e1_front_e2==-1:
#         continue
#     for e in entity2mask:
#         item[1] = item[1].replace(e,entity2mask[e])
#     new_sentence = replace_number_and_date_after_entity_removed(item[1])
#     item[1] = new_sentence

with open("./train_nlpcc6.csv",mode="w") as f,open("./test_nlpcc6.csv",mode="w") as f2:
    f.write("id")
    f.write(",")
    f.write("text")
    f.write(",")
    f.write("e1_front_in_sentence")
    f.write(",")
    f.write("e2_front_in_sentence")
    f.write(",")
    f.write("\n")

    f2.write("id")
    f2.write(",")
    f2.write("text")
    f2.write(",")
    f2.write("e1_front_in_sentence")
    f2.write(",")
    f2.write("e2_front_in_sentence")
    f2.write(",")
    f2.write("\n")
    for index,item in enumerate(train_data):
        # if random.randint(1,10)!=1:
        if True:
            if item[2]==-1:
                continue
            f.write("-1")
            f.write(",")
            f.write(item[1])
            f.write(",")
            temp = [0]*2
            if item[2]==True:
                temp[0]=1
            else:
                temp[1]=1
            for i in temp:
                f.write(str(i))
                f.write(",")
            f.write("\n")
        else:
            f2.write("-1")
            f2.write(",")
            f2.write(item[1])
            f2.write(",")
            temp = [0] * 2
            if item[2] == True:
                temp[0] = 1
            else:
                temp[1] = 1
            for i in temp:
                f2.write(str(i))
                f2.write(",")
            f2.write("\n")


